use crate::std_extensions::primitive::ExtendStr;
use std::collections::HashMap;
use std::rc::Rc;

#[derive(Clone, Debug)]
pub enum Trie<T> {
    Leaf,
    Node {
        key: String,
        val: Option<T>,
        children: HashMap<char, Rc<Trie<T>>>,
    },
}

use self::Trie::{Leaf, Node};

pub trait Map<K, V> {
    fn empty() -> Self;
    fn insert(&self, _: K, _: V) -> Self;
    fn find(&self, _: K) -> V;
}

impl<T: Clone> Map<String, T> for Trie<T> {
    fn empty() -> Trie<T> {
        Leaf
    }

    fn insert(&self, k: String, v: T) -> Trie<T> {
        let add = |t: &Trie<T>, k: String, v: T| -> Trie<T> {
            match t {
                Leaf => panic!("undefined"),
                Node {
                    key: _,
                    val: _,
                    children,
                } => match children.get(&k.char_at(0)) {
                    Some(n) => n.insert(k, v),
                    None => Node {
                        key: k,
                        val: Some(v),
                        children: HashMap::new(),
                    },
                },
            }
        };
        match self {
            Leaf => Node {
                key: k,
                val: Some(v),
                children: HashMap::new(),
            },
            Node { key, val, children } => {
                let i = k.chars().zip(key.chars()).take_while(|t| t.0 == t.1).count();
                if i == k.len() && i == key.len() {
                    Node {
                        key: k,
                        val: Some(v),
                        children: children.clone(),
                    }
                } else if i == key.len() {
                    let k1 = &k[i..];
                    let mut children = children.clone();
                    children.insert(k1.char_at(0), Rc::new(add(self, k1.to_string(), v)));
                    Node {
                        key: key.clone(),
                        val: val.clone(),
                        children,
                    }
                } else if i == k.len() {
                    let k1 = &key[i..];
                    let mut tmp_map = HashMap::new();
                    tmp_map.insert(
                        k1.char_at(0),
                        Rc::new(Node {
                            key: k1.to_string(),
                            val: val.clone(),
                            children: children.clone(),
                        }),
                    );
                    Node {
                        key: k,
                        val: Some(v),
                        children: tmp_map,
                    }
                } else {
                    let common = &k[..i];
                    let k1 = &key[i..];
                    let k2 = &k[i..];
                    let mut tmp_map = HashMap::new();
                    tmp_map.insert(
                        k1.char_at(0),
                        Rc::new(Node {
                            key: k1.to_string(),
                            val: val.clone(),
                            children: children.clone(),
                        }),
                    );
                    tmp_map.insert(
                        k2.char_at(0),
                        Rc::new(Node {
                            key: k2.to_string(),
                            val: Some(v),
                            children: HashMap::new(),
                        }),
                    );

                    Node {
                        key: common.to_string(),
                        val: None,
                        children: tmp_map,
                    }
                }
            }
        }
    }

    fn find(&self, k: String) -> T {
        match self {
            Leaf => panic!("empty tree"),
            Node { key, val, children: _ } if k == *key => match val {
                Some(v) => v.clone(),
                None => panic!("doesn't exist"),
            },
            Node { key, val: _, children } if k.starts_with(key) => match children.get(&k.char_at(key.len())) {
                Some(t) => t.find(k[key.len()..].to_string()),
                None => panic!("doesn't exist"),
            },
            _ => panic!("doesn't exist"),
        }
    }
}

#[test]
fn trie_test() {
    let t: Trie<usize> = Map::empty();
    let t2 = t.insert("test".to_string(), 0).insert("two".to_string(), 1);
    assert_eq!(t2.find("test".to_string()), 0);
    assert_eq!(t2.find("two".to_string()), 1);
}
